import re
from bs4 import BeautifulSoup
import logging

def remove_body_html_tags(html_string: str) -> str:
    """
    删除 HTML 字符串中的 </body> 和 </html> 标签

    Args:
        html_string (str): 输入的 HTML 字符串

    Returns:
        str: 删除后的 HTML 字符串
    """
    # 使用正则表达式删除 </body> 和 </html> 标签
    return re.sub(r'</?body>|</?html>', '', html_string)

def extract_first_rowspan(html_string: str) -> int:
    """
    提取HTML表格中第一个<tr>标签中第一个<td>标签的rowspan属性值

    Args:
        html_string (str): 输入的HTML表格字符串

    Returns:
        int: 第一个<td>标签的rowspan值，若没有rowspan属性则返回1
    """
    # 使用BeautifulSoup解析HTML字符串
    soup = BeautifulSoup(html_string, 'html.parser')

    # 获取第一个<tr>标签
    first_tr = soup.find_all('tr')[0]

    # 查找第一个<td>标签，并获取它的rowspan属性
    first_td = first_tr.find('td')

    # 获取rowspan属性，如果不存在则返回1
    rowspan = first_td.get('rowspan', 1)

    return int(rowspan)  # 返回rowspan的数值，默认为1

def get_row_colspan(html: str, row_index: int) -> int:
    """
    获取HTML表格中指定行的所有<td>和<th>的colspan总数。
    使用BeautifulSoup解析HTML，找到指定<tr>中的所有<td>和<th>元素，
    累加它们的colspan值（默认为1）。
    """
    soup = BeautifulSoup(html, 'html.parser')
    rows = soup.find_all('tr')
    
    if -len(rows) <= row_index < len(rows):
        cols = rows[row_index].find_all(['td', 'th'])
        return sum(int(col.get('colspan', 1)) for col in cols)
    
    return 0  # 如果索引超出范围，返回0

def get_elements_by_id(html: str, element_id: str):
    """
    获取HTML表格中指定id的<td>和<th>元素。
    使用BeautifulSoup解析HTML，找到匹配指定id的<td>和<th>元素，
    返回它们的列表。
    """
    soup = BeautifulSoup(html, 'html.parser')
    elements = soup.find_all(['td', 'th'], id=element_id)
    return elements


def row_merge(table1, table2):
    try:
        # 使用BeautifulSoup解析HTML
        soup1 = BeautifulSoup(table1, 'html.parser')
        soup2 = BeautifulSoup(table2, 'html.parser')


        # 获取第一个<tr>的所有<td>标签
        table2_tr = soup2.find_all('tr')[0]
        td_tags = table2_tr.find_all('td')
        table2_td_count = len(table2_tr.find_all('td'))

        # 遍历第一个<tr>中的所有<td>标签，查找为空的<td>并处理
        for i, td in enumerate(table2_tr.find_all('td')):
            if not td.text.strip():  # 如果<td>内容为空
                # 统计table2首行每个<td>的colspan属性，如果没有colspan则为1
                table2_colspan = sum(int(td.get('colspan', 1)) for td in table2_tr.find_all('td'))
                # 获取当前td对的colspan和rowspan属性，如果没有则默认为1
                colspan_value = int(td.get('colspan', 1))
                if colspan_value != 1:
                    break
                rowspan_value = int(td.get('rowspan', 1))
                # 在table1中找到最后一个包含colspan_value的<tr>
                current_rowspan = 0
                for tr in reversed(soup1.find_all('tr')):  # 从最后一个<tr>开始查找
                    current_rowspan += 1
                    table1_colspan = sum(
                        int(td1.get('colspan', 1)) for td1 in tr.find_all('td')
                    )
                    if table1_colspan == table2_colspan:  # 找到最后一个包含colspan_value的<tr>
                        id_colspan = []
                        for td1 in tr.find_all('td'):
                            if sum(id_colspan) + int(td1.get('colspan', 1)) <= colspan_value:
                                id_colspan.append(int(td1.get('colspan', 1)))
                            else:
                                break
                        if sum(id_colspan) == int(colspan_value):
                            for td_id in range(len(id_colspan)):
                                first_td_in_tr = tr.find_all('td')[td_id]
                                new_rowspan = current_rowspan + rowspan_value
                                first_td_in_tr['rowspan'] = str(new_rowspan)
                            # 删除content2中的对应<td>
                            td.decompose()

                        break
                    elif table1_colspan > table2_colspan:
                        break
            else:
                break

        # 返回修改后的content1和content2
        return str(soup1), str(soup2)
    except Exception as e:
        logging.error(f"Error in row_merge: {e}")
        logging.info(f"Table1: {repr(table1)[:3000]}")
        logging.info(f"Table2: {repr(table2)[:3000]}")
        return table1, table2


def row_merge1(table1, table2):
    # 使用BeautifulSoup解析HTML
    soup1 = BeautifulSoup(table1, 'html.parser')
    soup2 = BeautifulSoup(table2, 'html.parser')

    # 获取table2的首行tr对table2_tr，计算出其整行colspan值的和为colspan2
    
    # 如果table2_tr 的第一个td对的内容为空
    ## 删除table2的首行的第一个td对，返回其colspan值和rowspan值
    ## 并在table1进行遍历，找出与colspan2相等的tr对，并在其对应的td对（colspan值和要与下面的相等）中，增加对应的rowspan值
    # 循环上面的过程，直至table2的首行tr的第一个td不为空 
    
    
    # 获取第一个<tr>的所有<td>标签
    first_tr = soup2.find_all('tr')[0]
    td_tags = first_tr.find_all('td')
    first_td_count = len(first_tr.find_all('td'))

    # 统计每个<td>的colspan属性，如果没有colspan则为1
    colspan_count = sum(int(td.get('colspan', 1)) for td in td_tags)

    # 遍历第一个<tr>中的所有<td>标签，查找为空的<td>并处理
    for i, td in enumerate(first_tr.find_all('td')):
        colspan_value = td.get('colspan',1)
        if not td.text.strip() and  colspan_value == 1:  # 判断<td>是否为空
            rowspan_value = td.get('rowspan', 1)  # 获取rowspan属性，默认为1
            # 在content1中找到最后一个包含n-i+1个<td>的<tr>
            current_rowspan = 0

            for tr in reversed(soup1.find_all('tr')):  # 从最后一个<tr>开始查找
                current_rowspan += 1
                td_colspan_sum = sum(
                    int(td1.get('colspan', 1)) for td1 in tr.find_all('td')
                )
                if td_colspan_sum >= colspan_count:  # 找到最后一个包含n-i+1个<td>的<tr>
                    # id = len(tr.find_all('td')) - len(first_tr.find_all('td'))
                    if len(tr.find_all('td')) < first_td_count:
                        id = first_td_count - len(first_tr.find_all('td'))
                    else:
                        id = len(tr.find_all('td')) - len(first_tr.find_all('td'))
                    first_td_in_tr = tr.find_all('td')[id]
                    new_rowspan = current_rowspan + int(rowspan_value)
                    first_td_in_tr['rowspan'] = str(new_rowspan)
                    # 删除content2中的对应<td>
                    td.decompose()
                    # 更新 td_count
                    colspan_count = sum(int(td.get('colspan', 1)) for td in first_tr.find_all('td'))
                    break
        else:
            break

    # 返回修改后的content1和content2
    return str(soup1), str(soup2)

def get_first_row_tables(html_string):
    """
    解析HTML字符串，获取第一个<tr>的所有<td>标签的文本内容，去除特殊字符和所有空白字符。

    参数:
        html_string (str): 要解析的HTML字符串

    返回:
        list: 清理后的文本内容列表
    """
    # 使用BeautifulSoup解析HTML字符串
    soup = BeautifulSoup(html_string, 'html.parser')

    # 获取第一个<tr>的所有<td>标签的文本内容，去除特殊字符和所有空白字符
    first_tr = soup.find_all('tr')[0]  # 获取第一个<tr>
    td_contents = [
        re.sub(r'\s+', '', re.sub(r'[^\w\s]', '', td.get_text(strip=True)))  # 去除所有空白字符和特殊字符
        for td in first_tr.find_all('td')
    ]

    return ''.join(td_contents)

def process_table_colspan(first_element: str, second_element: str):
    """
    处理表格的colspan属性，根据第二个元素的colspan调整第一个元素的colspan。

    Args:
        first_element (str): 第一个表格元素的HTML字符串
        second_element (str): 第二个表格元素的HTML字符串
    """
    # 使用BeautifulSoup解析第一个表格
    first_soup = BeautifulSoup(first_element, 'html.parser')
    first_trs = first_soup.find_all('tr')  # 获取所有<tr>标签

    # 如果first_element中有超过一个<tr>标签，直接返回原始的first_element
    if len(first_trs) > 1:
        return first_element

    first_tr = first_trs[0]  # 获取第一个<tr>标签
    first_tds = first_tr.find_all('td')  # 获取所有<td>标签

    # 获取第一个元素中第一个<tr>所有<td>的colspan属性，并求和，默认值为1
    b1 = [int(td.get('colspan', 1)) for td in first_tds]
    b1_value = sum(b1)  # 计算colspan值的总和

    # 使用BeautifulSoup解析第二个表格
    second_soup = BeautifulSoup(second_element, 'html.parser')
    second_tr = second_soup.find_all('tr')[0]  # 获取第二个表格的第一个<tr>标签
    second_tds = second_tr.find_all('td')  # 获取第二个表格中的所有<td>标签

    # 获取第二个元素中第一个<tr>所有<td>的colspan属性，并求和，默认值为1
    b2 = [int(td.get('colspan', 1)) for td in second_tds]
    b2_value = sum(b2)  # 计算colspan值的总和

    # 如果b2 > b1，则进行处理
    if b2_value > b1_value:
        ratio = b2_value // b1_value  # 计算比例c1
        remainder = b2_value % b1_value

        # 处理第一个元素的colspan
        for idx, td in enumerate(first_tds):
            if idx == 0:
                continue  # 跳过第一个td
            if remainder == 0:
                td['colspan'] = ratio  # 设置colspan为c1
            else:
                td['colspan'] = ratio + 1  # 设置colspan为c1+1

    # 返回修改后的第一个元素的HTML
    return str(first_soup)

def is_header_only_table(html_table: str) -> bool:
    """
    判断HTML表格是否仅包含标题行。
    
    规则：
    - 检查表格第一行的所有单元格的 rowspan 值。
    - 如果最大 rowspan 值大于或等于表格总行数，则表格只包含标题行。
    
    参数:
    html_table (str): HTML 格式的表格字符串。
    
    返回:
    bool: 如果表格只有标题行，则返回 True，否则返回 False。
    """
    soup = BeautifulSoup(html_table, 'html.parser')
    table = soup.find('table')
    if not table:
        return False
    
    rows = table.find_all('tr')
    if not rows:
        return False
    
    total_rows = len(rows)
    first_row = rows[0]
    max_rowspan = max(
        int(cell.get('rowspan', 1)) for cell in first_row.find_all(['td', 'th'])
    )
    
    return max_rowspan >= total_rows

def table_merge(table1,table2):
    
    table1_colspan = get_row_colspan(table1,-1)
    table2_colspan = get_row_colspan(table2,0)
    table1_td = get_elements_by_id(table1,-1)
    table2_td = get_elements_by_id(table2,0)
    # 如果table1的最后一行的colspan数和table2的第一行的colspan数量相同，或者td对相同
    if table1_colspan == table2_colspan or table1_td == table2_td:
        table1,table2 = row_merge(table1, table2)
        table1 = '<table>' + re.sub(r'</?table>', '', table1+table2) + '</table>'
        table2 = None
    # 否则
    return table1,table2

def process_table_content(match):
    # 在这里处理匹配到的表格内容
    # 比如返回修改后的表格内容，或者直接返回原内容

    # 提取第一个表的字段行内容
    pattern = r'<tr>.*?</tr>'
    first_label_tr = re.search(pattern, match.group(0))
    ## 对字段行进行清理：匹配中文、英文、数字，去除特殊字符和空白
    first_label_tr_clear = get_first_row_tables(first_label_tr.group(0))
    # 划分多表
    tables = re.split(r'\n\s+', match.group(0))
    # 对每个表进行处理
    # 如果第一个表格只有标题行，则进行colspan值匹配
    if is_header_only_table(tables[0]):
        tables[0] = process_table_colspan(tables[0], tables[1])

    tables[0] = re.sub(r'</?table>', '', tables[0])
    # 进行重复标题行清理
    for i in range(1,len(tables)):
        # 提取每个表首行
        label_tr = re.search(pattern, tables[i])
        label_tr_clear = get_first_row_tables(label_tr.group(0))
        ## 多字段处理
        if first_label_tr_clear == label_tr_clear:
            # 判断字段行一共有多少行
            field_num = extract_first_rowspan(tables[i])
            pattern_field = pattern * field_num
            label_field = re.search(pattern_field, tables[i])
            tables[i] = tables[i].replace(label_field.group(0), "")
        tables[i] = re.sub(r'</?table>', '', tables[i])
    # 进行跨表行合并

    new_tables = [tables[0]]
    for i in range(len(tables)-1):
        table1,table2 = table_merge(new_tables[-1],tables[i+1])
        if table1:
            new_tables[-1] = table1
        if table2:
            new_tables.append(table2)

    # end_label = '<table>'+''.join(tables) + '</table>'
    end_tables = ['<table>'+ table + '</table>' for table in new_tables]
    end_tables = '\n'.join(end_tables)
    return end_tables  # 这里返回原内容


def tables_merge(text):
    # 正则表达式：匹配跨页表格的所以内容
    table_pattern = r'<table>.*?</table>(\s*<table>.*?</table>)+'


    # 使用re.sub替换匹配的部分，使用process_table_content处理每个匹配到的表格
    result = re.sub(table_pattern, process_table_content, text)

    return result

def table_md_change(md_content):

    # 删除body和html标签
    md_content = remove_body_html_tags(md_content)

    # 进行跨页表格合并
    md_content = tables_merge(md_content)
    return md_content

if __name__ == "__main__":
    logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s"
)
    # 读取文件内容
    file_path = '/home/fengjie/doc-parser/MinerU/output/√3/√3_pre.md'  # 替换为你的文件路径
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()

    result = table_md_change(content)
    file_path2 = '/home/fengjie/doc-parser/MinerU/output/√3/√3.md'  # 替换为你想要保存的文件路径
    with open('your_file.md', 'w', encoding='utf-8') as f:
        f.write(result)
    